#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from mitmproxy.options import Options
from mitmproxy.tools.dump import DumpMaster
from mitmproxy import http
import yaml
import logging
import colorlog
import os
import asyncio
import re
import json
import time

def configure_logging(log_level='INFO'):
    """配置日志系统"""
    # 将字符串级别转换为logging常量
    level_map = {
        'DEBUG': logging.DEBUG,
        'INFO': logging.INFO, 
        'WARNING': logging.WARNING,
        'ERROR': logging.ERROR,
        'CRITICAL': logging.CRITICAL
    }
    
    # 验证日志级别
    if log_level.upper() not in level_map:
        print(f"警告: 无效的日志级别 '{log_level}'，使用默认级别 INFO")
        log_level = 'INFO'
    
    # 配置日志处理器
    handler = colorlog.StreamHandler()
    handler.setFormatter(colorlog.ColoredFormatter(
        '%(log_color)s%(asctime)s %(levelname)s: %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S',
        log_colors={
            'DEBUG': 'cyan',
            'INFO': 'green',
            'WARNING': 'yellow',
            'ERROR': 'red',
            'CRITICAL': 'red,bg_white',
        }
    ))
    
    logger = logging.getLogger('proxy')
    # 清除现有的处理器
    logger.handlers.clear()
    logger.addHandler(handler)
    logger.setLevel(level_map[log_level.upper()])
    
    return logger

# 默认日志配置
logger = configure_logging()

class ProxyAddon:
    def __init__(self, config_path='config.yaml', log_level=None):
        self.config = self.load_config(config_path)
        self.rules = self.config.get('rules', [])
        # 获取默认目标端口配置
        self.default_target_port = self.config.get('server', {}).get('proxy', {}).get('default_target_port', 80)
        
        # 正则表达式编译缓存
        self.regex_cache = {}
        self._compile_regex_patterns()
        
        # 配置日志级别 (CLI参数 > 配置文件 > 默认值)
        if log_level:
            # CLI参数优先
            effective_log_level = log_level
        else:
            # 使用配置文件中的级别
            effective_log_level = self.config.get('logging', {}).get('level', 'INFO')
        
        # 重新配置日志系统
        global logger
        logger = configure_logging(effective_log_level)
        
        logger.info(f"日志级别设置为: {effective_log_level}")
        
        # 统计规则类型和状态
        proxy_rules = 0
        mock_rules = 0
        enabled_rules = 0
        disabled_rules = 0
        
        for rule in self.rules:
            if rule.get('enabled', True):
                enabled_rules += 1
                if self.is_mock_rule(rule):
                    mock_rules += 1
                else:
                    proxy_rules += 1
            else:
                disabled_rules += 1
        
        logger.info(f"加载了 {len(self.rules)} 条规则: {proxy_rules} 条代理规则, {mock_rules} 条Mock规则")
        logger.info(f"规则状态: {enabled_rules} 条启用, {disabled_rules} 条禁用")

    def _compile_regex_patterns(self):
        """预编译所有正则表达式模式"""
        regex_count = 0
        for i, rule in enumerate(self.rules):
            pattern_type = rule.get('pattern_type', 'string')
            if pattern_type == 'regex':
                pattern = rule.get('pattern', '')
                try:
                    compiled_regex = re.compile(pattern)
                    self.regex_cache[i] = compiled_regex
                    regex_count += 1
                    logger.debug(f"编译正则表达式规则 {i}: {pattern}")
                except re.error as e:
                    logger.error(f"正则表达式编译失败，规则 {i} ({pattern}): {e}")
                    # 将无效的正则规则降级为字符串匹配
                    rule['pattern_type'] = 'string'
                    logger.warning(f"规则 {i} 已降级为字符串匹配模式")
        
        if regex_count > 0:
            logger.info(f"预编译了 {regex_count} 条正则表达式规则")

    def load_config(self, config_path):
        """加载配置文件"""
        if not os.path.exists(config_path):
            logger.warning(f"配置文件 {config_path} 不存在，使用默认配置")
            return {'rules': [], 'logging': {'level': 'INFO'}}
        
        with open(config_path, 'r') as f:
            config = yaml.safe_load(f)
            # 确保logging配置存在，提供默认值
            if 'logging' not in config:
                config['logging'] = {'level': 'INFO'}
            
            # 处理规则的向后兼容性和Mock配置验证
            if 'rules' in config:
                for rule in config['rules']:
                    self._validate_rule(rule)
            
            return config
    
    def _validate_rule(self, rule):
        """验证规则配置"""
        # 设置默认的enabled状态
        if 'enabled' not in rule:
            rule['enabled'] = True  # 默认启用
        
        # 验证enabled字段
        if not isinstance(rule['enabled'], bool):
            logger.warning(f"规则 {rule.get('pattern', 'unknown')} 的enabled字段应为布尔值，已设为true")
            rule['enabled'] = True
        
        # 验证pattern_type字段
        if 'pattern_type' not in rule:
            rule['pattern_type'] = 'string'  # 默认为字符串匹配
        
        pattern_type = rule['pattern_type']
        if pattern_type not in ['string', 'regex']:
            logger.warning(f"规则 {rule.get('pattern', 'unknown')} 的pattern_type无效: {pattern_type}，已设为string")
            rule['pattern_type'] = 'string'
        
        # 验证正则表达式语法
        if pattern_type == 'regex':
            pattern = rule.get('pattern', '')
            try:
                re.compile(pattern)
            except re.error as e:
                logger.error(f"规则 {rule.get('pattern', 'unknown')} 的正则表达式语法错误: {e}")
                # 降级为字符串匹配
                rule['pattern_type'] = 'string'
                logger.warning(f"规则已降级为字符串匹配模式")
        
        if 'mock_response' in rule:
            mock_config = rule['mock_response']
            
            # 验证状态码
            if 'status' in mock_config:
                status = mock_config['status']
                if not isinstance(status, int) or status < 100 or status > 599:
                    logger.warning(f"规则 {rule.get('pattern', 'unknown')} 的状态码无效: {status}")
                    mock_config['status'] = 200
            else:
                mock_config['status'] = 200
            
            # 验证内容配置
            if 'content' in mock_config and 'content_file' in mock_config:
                logger.warning(f"规则 {rule.get('pattern', 'unknown')} 同时配置了content和content_file，将使用content_file")
                del mock_config['content']
            
            # 验证外部文件是否存在
            if 'content_file' in mock_config:
                file_path = mock_config['content_file']
                if not os.path.exists(file_path):
                    logger.error(f"Mock文件不存在: {file_path}")
                    # 提供默认内容
                    mock_config['content'] = '{"error": "Mock file not found"}'
                    mock_config['status'] = 500
                    del mock_config['content_file']

    def should_proxy(self, flow):
        """判断请求是否需要代理或Mock"""
        url = flow.request.pretty_url
        
        for i, rule in enumerate(self.rules):
            # 检查规则是否启用
            if not rule.get('enabled', True):
                logger.debug(f"跳过禁用规则: {rule['pattern']} - {url}")
                continue
            
            # 根据匹配模式进行匹配
            pattern_type = rule.get('pattern_type', 'string')
            pattern = rule.get('pattern', '')
            
            matched = False
            
            if pattern_type == 'regex':
                # 正则表达式匹配
                if i in self.regex_cache:
                    try:
                        matched = bool(self.regex_cache[i].search(url))
                        if matched:
                            logger.debug(f"正则匹配成功: {pattern} -> {url}")
                    except Exception as e:
                        logger.error(f"正则匹配执行失败: {e}")
                        # 降级到字符串匹配
                        matched = pattern in url
                        logger.debug(f"降级为字符串匹配: {pattern} in {url} = {matched}")
                else:
                    # 缓存中没有编译的正则，降级为字符串匹配
                    matched = pattern in url
                    logger.debug(f"正则缓存缺失，使用字符串匹配: {pattern} in {url} = {matched}")
            else:
                # 字符串包含匹配（默认行为）
                matched = pattern in url
                logger.debug(f"字符串匹配: {pattern} in {url} = {matched}")
            
            if matched:
                logger.debug(f"规则匹配: {pattern} ({pattern_type}) -> {url}")
                return rule
        
        return None
    
    def is_mock_rule(self, rule):
        """判断规则是否为Mock规则"""
        if not rule:
            return False
        # 如果有mock_response字段，就是Mock规则
        if 'mock_response' in rule:
            return True
        # 如果type明确指定为mock，也是Mock规则
        if rule.get('type') == 'mock':
            return True
        return False
    
    def load_mock_content(self, mock_config):
        """加载Mock内容"""
        if 'content_file' in mock_config:
            try:
                with open(mock_config['content_file'], 'r', encoding='utf-8') as f:
                    return f.read()
            except Exception as e:
                logger.error(f"读取Mock文件失败: {e}")
                return '{"error": "Failed to load mock content"}'
        elif 'content' in mock_config:
            content = mock_config['content']
            # 如果content是字典或列表，转换为JSON字符串
            if isinstance(content, (dict, list)):
                return json.dumps(content, ensure_ascii=False, indent=2)
            return str(content)
        else:
            return '{"message": "Mock response"}'
    
    def create_mock_response(self, mock_config):
        """创建Mock响应"""
        # 加载内容
        content = self.load_mock_content(mock_config)
        
        # 准备响应头
        headers = {}
        
        # 设置Content-Type
        if 'content_type' in mock_config:
            headers['Content-Type'] = mock_config['content_type']
        else:
            # 自动推断内容类型
            if mock_config.get('content_file', '').endswith('.json') or content.strip().startswith(('{', '[')):
                headers['Content-Type'] = 'application/json; charset=utf-8'
            elif mock_config.get('content_file', '').endswith('.html') or content.strip().startswith('<'):
                headers['Content-Type'] = 'text/html; charset=utf-8'
            else:
                headers['Content-Type'] = 'text/plain; charset=utf-8'
        
        # 添加自定义头部
        if 'headers' in mock_config:
            headers.update(mock_config['headers'])
        
        # 添加Mock标识头部
        headers['X-Mock-Response'] = 'true'
        
        # 获取状态码
        status_code = mock_config.get('status', 200)
        
        # 处理延迟
        if 'delay' in mock_config:
            delay_ms = mock_config['delay']
            if isinstance(delay_ms, (int, float)) and delay_ms > 0:
                time.sleep(delay_ms / 1000.0)  # 转换为秒
        
        # 创建响应
        return http.Response.make(
            status_code,
            content.encode('utf-8'),
            headers
        )

    def transform_path(self, original_path, rule):
        """转换请求路径"""
        if 'strip_prefix' in rule:
            return re.sub(f"^{rule['strip_prefix']}", "", original_path)
        return original_path

    async def request(self, flow):
        """处理请求"""
        rule = self.should_proxy(flow)
        if rule:
            original_url = flow.request.pretty_url
            
            # 检查是否为Mock规则
            if self.is_mock_rule(rule):
                # 处理Mock响应
                mock_config = rule.get('mock_response', {})
                if not mock_config and rule.get('type') == 'mock':
                    # type为mock但没有mock_response配置，返回404
                    flow.response = http.Response.make(
                        404,
                        '{"error": "Mock configuration not found"}',
                        {"Content-Type": "application/json", "X-Mock-Response": "true"}
                    )
                    logger.warning(f"Mock规则配置缺失: {original_url}")
                    return
                
                try:
                    flow.response = self.create_mock_response(mock_config)
                    logger.info(f"Mock响应: {mock_config.get('status', 200)} - {original_url}")
                    logger.debug(f"Mock详情: {flow.request.method} {original_url} -> {len(flow.response.content)}bytes")
                except Exception as e:
                    logger.error(f"创建Mock响应失败: {e}")
                    flow.response = http.Response.make(
                        500,
                        '{"error": "Failed to create mock response"}',
                        {"Content-Type": "application/json", "X-Mock-Response": "true"}
                    )
                return
            
            # 处理普通代理请求
            if 'target' in rule:
                target = rule['target']
                
                # 修改请求目标地址
                flow.request.scheme = target.get('scheme', 'http')
                flow.request.host = target.get('host', 'localhost')
                flow.request.port = target.get('port', self.default_target_port)
                
                # 转换路径
                flow.request.path = self.transform_path(flow.request.path, rule)
                
                logger.info(f"代理请求: {original_url} -> {flow.request.scheme}://{flow.request.host}:{flow.request.port}{flow.request.path}")
                logger.debug(f"详细请求信息: {flow.request.method} {original_url}")
            else:
                logger.warning(f"规则缺少target配置: {rule.get('pattern', 'unknown')}")

    async def response(self, flow):
        """处理响应"""
        # 将详细的响应信息设为DEBUG级别
        logger.debug(f"响应详情: {flow.response.status_code} {flow.request.method} {flow.request.pretty_url} ({len(flow.response.content)}bytes)")
        
        # 检查是否为Mock响应
        if flow.response and 'X-Mock-Response' in flow.response.headers:
            # Mock响应已经在request阶段记录过日志，这里不重复记录
            return
        
        # 只有在代理规则匹配时才记录INFO级别日志
        rule = self.should_proxy(flow)
        if rule and not self.is_mock_rule(rule):
            logger.info(f"代理响应: {flow.response.status_code} - {flow.request.pretty_url}")

async def start_proxy_async(host=None, port=None, config_path='config.yaml', log_level=None):
    """异步启动代理服务器"""
    # 如果没有指定host和port，从配置文件读取
    if host is None or port is None:
        try:
            with open(config_path, 'r') as f:
                config = yaml.safe_load(f)
                server_config = config.get('server', {}).get('proxy', {})
                if host is None:
                    host = server_config.get('host', '0.0.0.0')
                if port is None:
                    port = server_config.get('port', 8080)
        except:
            # 配置文件读取失败，使用默认值
            if host is None:
                host = '0.0.0.0'
            if port is None:
                port = 8080
    
    logger.info(f"启动代理服务器 {host}:{port}")
    
    # 自动更新系统代理设置
    await auto_update_system_proxy(config_path)
    
    opts = Options(
        listen_host=host,
        listen_port=port,
    )
    
    master = DumpMaster(opts)
    master.addons.add(ProxyAddon(config_path, log_level))
    
    await master.run()
    return master

async def auto_update_system_proxy(config_path='config.yaml'):
    """自动更新系统代理设置"""
    try:
        # 导入SystemProxyManager（延迟导入避免循环依赖）
        import sys
        import os
        sys.path.append(os.path.dirname(os.path.abspath(__file__)))
        
        from web_manager.proxy.system_manager import SystemProxyManager
        
        system_manager = SystemProxyManager(config_path)
        success = system_manager.enable_proxy()
        
        if success:
            logger.info("自动更新系统代理设置成功")
        else:
            logger.warning("自动更新系统代理设置失败，请手动设置")
            
    except Exception as e:
        logger.warning(f"自动更新系统代理时出错: {e}")
        logger.info("代理服务器正常启动，请手动设置系统代理")

def start_proxy(host=None, port=None, config_path='config.yaml', log_level=None):
    """启动代理服务器"""
    loop = asyncio.get_event_loop()
    master = None
    
    try:
        master = loop.run_until_complete(
            start_proxy_async(host, port, config_path, log_level)
        )
    except KeyboardInterrupt:
        logger.info("正在停止代理服务器...")
    finally:
        if master:
            master.shutdown()

if __name__ == '__main__':
    import argparse
    
    parser = argparse.ArgumentParser(description='启动代理服务器')
    parser.add_argument('--host', help='监听地址 (默认从配置文件读取，如无配置则为0.0.0.0)')
    parser.add_argument('--port', type=int, help='监听端口 (默认从配置文件读取，如无配置则为8080)')
    parser.add_argument('--config', default='config.yaml', help='配置文件路径')
    parser.add_argument('--log-level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], 
                       help='日志级别 (覆盖配置文件设置)')
    
    args = parser.parse_args()
    
    start_proxy(args.host, args.port, args.config, args.log_level) 